import argparse
import os
import numpy as np
import scipy.interpolate as interp
from scipy.interpolate import Rbf
import time

def nearest_neighbor_interpolation(data, r,z):
    """
    使用临近插值法对二维数据进行插值
    :param data: 原始数据，格式为 [(x1, y1, z1), (x2, y2, z2), ...]
    :param new_points: 需要插值的点，格式为 [(x, y), (x, y), ...]
    :return: 插值后的z值列表
    """
    data = np.array(data)
    x_values = data[:, 0]
    y_values = data[:, 1]
    z_values = data[:, 2]
    # distances = np.sqrt((x_values - pts[0])**2 + (y_values - pts[1])**2)
    # nearest_index = np.argmin(distances)
    # interpolated_z=z_values[nearest_index]

    # rbf = Rbf(x_values, y_values, z_values, function='multiquadric')

    interpolated_z = interp.griddata(data[:, :2], data[:, 2], (r, z), method='linear')

    return interpolated_z

def calculate_element_centroids(nodes, element):
    """
    计算每个单元的中心点坐标
    :param nodes: 节点信息字典，格式为 {节点ID: [x, y, z]}
    :param elements: 单元信息字典，格式为 {单元ID: {'type': 单元类型, 'nodes': [节点ID, ...]}}
    :return: 单元中心点坐标字典，格式为 {单元ID: [x, y, z]}
    """

    node_ids = element['nodes']
    node_coords = np.array([nodes[node_id] for node_id in node_ids])
    centroid = np.mean(node_coords, axis=0)
    
    return centroid.tolist()

def cartesian_to_cylindrical(cartesian_coords):
    """
    将直角坐标转换为柱坐标
    :param cartesian_coords: 直角坐标列表，格式为 [(x1, y1, z1), (x2, y2, z2), ...]
    :return: 柱坐标列表，格式为 [(r1, theta1, z1), (r2, theta2, z2), ...]
    """
    cylindrical_coords = []
    
    for (x, y, z) in cartesian_coords:
        r = np.sqrt(x**2 + y**2)
        cylindrical_coords.append((r, z))
    
    return cylindrical_coords

def read_calculix_inp(file_path):
    """
    读取CalculiX输入文件(.inp)并解析数据
    :param file_path: .inp文件的路径
    :return: 包含节点、单元和单元集合信息的字典
    """
    nodes = {}
    elements = {}
    element_sets = {}
    
    with open(file_path, 'r',encoding='utf-8') as file:
        lines = file.readlines()
        
        reading_nodes = False
        reading_elements = False
        reading_element_set = False
        
        element_type = None
        current_element_set = None
        
        for line in lines:
            line = line.strip()
            
            if line.startswith('*'):
                reading_nodes = False
                reading_elements = False
                reading_element_set = False
                
                if line.startswith('*Node'):
                    reading_nodes = True
                    continue
                elif line.startswith('*Element'):
                    reading_elements = True
                    parts = line.split(',')
                    if len(parts) > 1:
                        element_type = parts[1].split('=')[1].strip()
                    if len(parts)>2:
                        reading_element_set=True
                        current_element_set = parts[2].split('=')[1].strip()
                        element_sets[current_element_set] = []
                    else:
                        element_type = 'Unknown'
                    continue
                elif line.startswith('*Nset'):
                    break
            
            if reading_nodes and line:
                parts = line.split(',')
                node_id = int(parts[0].strip())
                coords = [float(coord.strip()) for coord in parts[1:]]
                nodes[node_id] = coords
            
            if reading_elements and line:
                parts = line.split(',')
                elem_id = int(parts[0].strip())
                node_ids = [int(node_id.strip()) for node_id in parts[1:]]
                elements[elem_id] = {
                    'type': element_type,
                    'nodes': node_ids
                }
                element_sets[current_element_set].append(elem_id)

            if reading_element_set and line:
                parts = line.split(',')
                for part in parts:
                    element_sets[current_element_set].append(int(part.strip()))
    
    return {
        'nodes': nodes,
        'elements': elements,
        'element_sets': element_sets
    }

def get_mag_force(file_path):
    """
    读取comsol计算结果洛伦兹力结果数据
    :param vertices: 洛伦兹力结果文件路径
    :return: 洛伦兹力数据[(r,z,value),...(r,z,value)]
    """
    with open(file_path, 'r',encoding='utf-8') as file:
        lines = file.readlines()
        data = []
        for line in lines:
            if line.startswith("%"):
                continue
            row = line.strip().split()
            row = [float(val) for val in row if val.strip()]
            data.append(row)
        return data
    
def add_mag_force(inp,force_r,force_z):
    mag_force={}
    nodes=inp["nodes"]
    elements=inp["elements"]
    eleLst=[]
    centerLst=[]
    thetaLst=[]
    rLst=[]
    hLst=[]
    for set_name, eles in inp['element_sets'].items():
        set_id=int(set_name.split("-")[-1])
        if set_id<=206:
            for ele in eles:
                eleLst.append(ele)
                element=elements[ele]
                center=calculate_element_centroids(nodes,element)
                radius = np.sqrt(center[0]**2 + center[2]**2)
                theta = np.arctan2(center[2], center[0])
                height=center[1]
                centerLst.append(center)
                thetaLst.append(theta)
                rLst.append(radius)
                hLst.append(height)
        else:
            continue
    rRes=nearest_neighbor_interpolation(force_r,rLst,hLst)
    zRes=nearest_neighbor_interpolation(force_z,rLst,hLst)
    for i in range(len(eleLst)):
        mag_force[eleLst[i]]=[rRes[i]*np.cos(thetaLst[i]),zRes[i],rRes[i]*np.sin(thetaLst[i])]
    return mag_force

def read_file_content(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
    except FileNotFoundError:
        print("文件未找到")
    except:
        print("读取文件出错")
    else:
        return content

def delete_file(filename):
    if os.path.exists(filename):
        os.remove(filename)
        print(f"文件 {filename} 删除成功！")
    else:
        print(f"文件 {filename} 不存在。")

if __name__ == "__main__":
    # Command line arguments
    ap = argparse.ArgumentParser()
    ap.add_argument('filename', type=str, help='initial model file name with extension inp')
    ap.add_argument('force_r', type=str,  help='input comsol mag force r file with extension txt')
    ap.add_argument('force_z', type=str,  help='input comsol mag force z file with extension txt')
    ap.add_argument('output', type=str,  help='output model file with extension inp')
    args = ap.parse_args()

    # Check arguments
    assert os.path.isfile(args.filename), 'calculate model file does not exist.'
    assert os.path.isfile(args.force_r), 'input comsol mag force r file does not exist.'
    assert os.path.isfile(args.force_z), 'input comsol mag force z file does not exist.'

    first_time = time.time()
    start_time = time.time()
    print("*****************************************")
    print(f"加载原始计算模型数据中......")
    inp=read_calculix_inp(args.filename)
    print(f"原始模型文件 {args.filename} 数据加载成功！")
    end_time = time.time()
    print(f"执行耗时： {0:.6f} s！".format(end_time - start_time))

    print("*****************************************")
    print(f"加载径向电磁力数据中......")
    start_time = time.time()
    force_r=get_mag_force(args.force_r)
    print(f"径向电磁力 {args.force_r} 数据加载成功！")
    end_time = time.time()
    print(f"执行耗时： {0:.6f} s！".format(end_time - start_time))

    print("*****************************************")
    print(f"加载轴向电磁力数据中......")
    start_time = time.time()
    force_z=get_mag_force(args.force_z)
    print(f"轴向电磁力 {args.force_z} 数据加载成功！")
    end_time = time.time()
    print(f"执行耗时： {0:.6f} s！".format(end_time - start_time))

    print("*****************************************")
    print(f"电磁力数据插值计算中......（较为耗时，请耐心等待）")
    start_time = time.time()
    mag_force=add_mag_force(inp,force_r,force_z)
    print(f"电磁力数据插值计算完成！")
    end_time = time.time()
    print(f"执行耗时： {0:.6f} s！".format(end_time - start_time))

    print("*****************************************")
    delete_file('ForceTemp.txt')
    print(f"电磁力计算模型数据导出中......")
    start_time = time.time()
    with open('ForceTemp.txt', 'w') as file:
        for ele_id,value in mag_force.items():
            mag=np.sqrt(value[0]**2+value[1]**2+value[2]**2)
            if mag==0:
                continue
            x=value[0]/mag
            y=value[1]/mag
            z=value[2]/mag
            format="{0}, Grav, {1:f}, {2:f}, {3:f}, {4:f}\n"
            file.writelines(format.format(ele_id, mag/(8.9), x, y, z))
    MAG_TEXT=read_file_content('ForceTemp.txt')
    INP_TEXT=read_file_content(args.filename)
    with open(args.output, 'w', encoding='utf-8') as file:
        file.write(INP_TEXT.replace('**MagForceReplacement**',MAG_TEXT))
    print(f"加载电磁力的计算模型文件已成功到处至： {args.output} ")
    end_time = time.time()
    print(f"执行耗时： {0:.6f} s！".format(end_time - start_time))
    print("*****************************************")
    print(f"执行总耗时： {0:.6f} s！".format(end_time - first_time))